# ai_search_bot.py
import os
import json
import urllib.parse
import json5
from typing import List, Dict, Any, Optional
from qwen_agent.agents import Assistant
from qwen_agent.tools.base import BaseTool, register_tool
from qwen_agent.gui import WebUI
import warnings
warnings.filterwarnings("ignore")


@register_tool('my_image_gen')
class MyImageGen(BaseTool):
    """AI图像生成工具"""
    description = 'AI 绘画（图像生成）服务，输入文本描述，返回基于文本信息绘制的图像 URL。'
    parameters = [{
        'name': 'prompt',
        'type': 'string',
        'description': '期望的图像内容的详细描述',
        'required': True
    }]

    def call(self, params: str, **kwargs) -> str:
        prompt = json5.loads(params)['prompt']
        prompt = urllib.parse.quote(prompt)
        return json5.dumps(
            {'image_url': f'https://image.pollinations.ai/prompt/{prompt}'},
            ensure_ascii=False
        )


class AdaptiveRetrievalConfig:
    """
    自适应检索配置类
    根据文档数量自动选择合适的检索方式
    """
    def __init__(self, 
                 doc_threshold: int = 50,
                 use_es_retrieval: bool = False,
                 es_config: Optional[Dict] = None):
        """
        初始化配置
        :param doc_threshold: 文档数量阈值，超过此数量使用ES检索
        :param use_es_retrieval: 是否强制使用ES检索
        :param es_config: ES配置信息
        """
        self.doc_threshold = doc_threshold
        self.use_es_retrieval = use_es_retrieval
        self.es_config = es_config or {
            "host": "https://localhost",
            "port": 9200,
            "user": "elastic",
            "password": "ncUg27UP5aMSzSElVXAX",
            "index_name": "adaptive_search_docs_index"
        }
    
    def should_use_es_retrieval(self, file_count: int) -> bool:
        """
        判断是否应该使用ES检索
        :param file_count: 文档数量
        :return: 是否使用ES检索
        """
        if self.use_es_retrieval:
            return True
        return file_count > self.doc_threshold


def init_agent_service(retrieval_config: Optional[AdaptiveRetrievalConfig] = None):
    """
    初始化具备自适应检索能力的助手服务
    :param retrieval_config: 检索配置
    """
    # 默认配置
    if retrieval_config is None:
        retrieval_config = AdaptiveRetrievalConfig()
    
    # 步骤 1: LLM 配置
    llm_cfg = {
        'model': 'qwen-plus',#'qwen-max',
        'model_server': 'dashscope',
        'api_key': os.getenv('DASHSCOPE_API_KEY'),
        'generate_cfg': {
            'top_p': 0.8
        }
    }

    # 步骤 2: 获取文件列表
    file_dir = os.path.join(os.path.dirname(__file__), 'docs')
    files = []
    if os.path.exists(file_dir):
        for file in os.listdir(file_dir):
            file_path = os.path.join(file_dir, file)
            if os.path.isfile(file_path):
                files.append(file_path)
    file_count = len(files)
    print(f'知识库文件列表: {files}')
    print(f'文档数量: {file_count}')

    # 步骤 3: 根据文档数量选择检索方式
    use_es_retrieval = retrieval_config.should_use_es_retrieval(file_count)
    
    if use_es_retrieval:
        print("使用 Elasticsearch 检索 (es_retrieval)")
        # 使用 Elasticsearch 检索配置
        rag_cfg = {
            "rag_backend": "elasticsearch",
            "es": retrieval_config.es_config,
            "parser_page_size": 500
        }
        # 工具配置 - 使用my_image_gen和code_interpreter工具
        tools_cfg = ['my_image_gen', 'code_interpreter']
    else:
        print("使用默认检索 (qwen-agent retrieval)")
        # 不使用特殊RAG配置，但必须启用本地RAG功能
        rag_cfg = {
            "rag_backend": "local",  # 添加本地RAG配置
            "parser_page_size": 500
        }
        # 工具配置 - 使用my_image_gen和code_interpreter工具
        tools_cfg = ['my_image_gen', 'code_interpreter']
    
    # 添加 tavily-mcp 网络搜索工具（仅在需要时添加）
    # 修改：将网络搜索工具放在最后，作为备选方案
    print("use_es_retrieval:", use_es_retrieval)
    if use_es_retrieval:
        tools_cfg.append({
            "mcpServers": {
                "tavily-mcp": {
                    "command": "npx",
                    "args": ["-y", "tavily-mcp@0.1.4"],
                    "env": {
                        "TAVILY_API_KEY": os.getenv('TAVILY_API_KEY', "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")
                    },
                    "disabled": False,
                    "autoApprove": []
                }
            }
        })
    print(tools_cfg)
    # 步骤 4: 系统指令
    system_instruction = '''你是一个AI助手，专门用于文档检索和问答任务。
请根据用户的问题，优先利用检索工具从本地知识库中查找最相关的信息。
如果本地知识库没有相关信息，再使用 tavily_search 工具从互联网上搜索，并结合这些信息给出专业、准确的回答。
知识库中是以前做过的示例，如果有类似的需要进行逻辑上的参考。
示例：
问题: 雇主安心保的赔偿范围
答案:雇主安心保赔偿范围：
1.	死亡赔偿金
2.	伤残赔偿金
3.	医疗费用
4.	误工费用
5.	法律诉讼费用
具体投保方案和保费请以投保页面展示和计算结果为准，不同地区会有承保方案和价格差异。

当前可用的知识库包括：
- 平安商业综合责任保险
- 雇主责任险知识库
- 企业团体综合意外险
- 财产一切险知识库
- 装修保险知识库

回答时请遵循以下原则：
1. 优先使用本地知识库的信息
2. 如果使用了网络搜索信息，请明确说明
3. 回答应准确、简洁，避免冗余信息
4. 如果不确定答案，请说明不确定，不要编造信息'''


    # 步骤 5: 创建智能体实例
    bot = Assistant(
        name="保险助手",
        description="一个用于保险的助手，能够处理保险相关的问题。优先使用内部保险知识库。", 
        llm=llm_cfg,
        system_message=system_instruction,
        function_list=tools_cfg,
        files=files,
        rag_cfg=rag_cfg  # 确保在所有情况下都设置rag_cfg
    )
    return bot, use_es_retrieval


def main(doc_threshold: int = 50, force_es_retrieval: bool = False):
    """
    启动 Web 图形界面
    :param doc_threshold: 文档数量阈值
    :param force_es_retrieval: 是否强制使用ES检索
    """
    try:
        # 创建检索配置
        es_config = {
            "host": "https://localhost",
            "port": 9200,
            "user": "elastic",
            "password": "ncUg27UP5aMSzSElVXAX",
            "index_name": "adaptive_search_docs_index"
        }
        
        retrieval_config = AdaptiveRetrievalConfig(
            doc_threshold=doc_threshold,
            use_es_retrieval=force_es_retrieval,
            es_config=es_config
        )
        
        print("正在启动 AI 搜索助手 Web 界面...")
        bot, use_es_retrieval = init_agent_service(retrieval_config)
        
        retrieval_method = "Elasticsearch 检索" if use_es_retrieval else "默认检索"
        print(f"当前使用的检索方式: {retrieval_method}")
        
        chatbot_config = {
            'prompt.suggestions': [
                '介绍下雇主责任险',
                '雇主责任险和工伤保险有什么主要区别？',
                '介绍一下平安商业综合责任保险的保障范围。',
                '施工保主要适用于哪些场景？',
                '请比较不同保险产品的优缺点'
            ]
        }
        WebUI(bot, chatbot_config=chatbot_config).run()
    except Exception as e:
        print(f"启动 Web 界面失败: {e}")
        print("请检查网络连接、API Key 以及 Elasticsearch 服务是否正常运行。")


def main_cli(doc_threshold: int = 50, force_es_retrieval: bool = False):
    """
    启动命令行交互界面
    :param doc_threshold: 文档数量阈值
    :param force_es_retrieval: 是否强制使用ES检索
    """
    try:
        # 创建检索配置
        es_config = {
            "host": "https://localhost",
            "port": 9200,
            "user": "elastic",
            "password": "ncUg27UP5aMSzSElVXAX",
            "index_name": "adaptive_search_docs_index"
        }
        
        retrieval_config = AdaptiveRetrievalConfig(
            doc_threshold=doc_threshold,
            use_es_retrieval=force_es_retrieval,
            es_config=es_config
        )
        
        print("正在启动 AI 搜索助手 CLI 界面...")
        bot, use_es_retrieval = init_agent_service(retrieval_config)
        
        retrieval_method = "Elasticsearch 检索" if use_es_retrieval else "默认检索"
        print(f"当前使用的检索方式: {retrieval_method}")
        print("输入 'quit' 或 'exit' 退出程序")
        
        # 对话历史
        messages = []
        while True:
            try:
                query = input('\n请输入您的问题: ').strip()
                if query.lower() in ['quit', 'exit', '退出']:
                    print("再见！")
                    break
                
                if not query:
                    print("问题不能为空！")
                    continue
                
                messages.append({'role': 'user', 'content': query})
                print("正在处理您的问题...")
                
                # 流式输出回答
                response = []
                for response_chunk in bot.run(messages=messages):
                    if response_chunk and response_chunk[-1]['role'] == 'assistant':
                        assistant_message = response_chunk[-1]
                        print(assistant_message.get('content', ''), end='', flush=True)
                    response = response_chunk
                
                print()  # 换行
                messages.extend(response)
                
            except KeyboardInterrupt:
                print("\n程序被用户中断")
                break
            except Exception as e:
                print(f"\n处理问题时出错: {e}")
                
    except Exception as e:
        print(f"启动 CLI 界面失败: {e}")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='AI 搜索助手')
    parser.add_argument('--mode', choices=['gui', 'cli'], default='gui', 
                        help='运行模式: gui(图形界面) 或 cli(命令行界面)')
    parser.add_argument('--threshold', type=int, default=50,
                        help='文档数量阈值，超过此数量使用ES检索 (默认: 50)')
    parser.add_argument('--force-es', action='store_true',
                        help='强制使用 Elasticsearch 检索')
    
    args = parser.parse_args()
    
    if args.mode == 'gui':
        main(doc_threshold=args.threshold, force_es_retrieval=args.force_es)
    else:
        main_cli(doc_threshold=args.threshold, force_es_retrieval=args.force_es)